from gmssl.sm2 import CryptSM2
from gmssl import func
from datetime import datetime, timedelta
import os
import json
import sqlite3
from cryptography.fernet import Fernet

DATABASE = 'certificates.db'


class Certificate:
    def __init__(self, subject, issuer, serial_number, not_before, not_after, public_key, signature, cert_type='encryption', version='v1', sig_alg='sm3'):
        self.version = version
        self.sig_alg = sig_alg
        self.subject = subject
        self.issuer = issuer
        self.serial_number = serial_number
        self.not_before = not_before
        self.not_after = not_after
        self.public_key = public_key
        self.signature = signature
        self.cert_type = cert_type

    def to_dict(self):
        return {
            'version': self.version,
            'sig_alg': self.sig_alg,
            'subject': self.subject,
            'issuer': self.issuer,
            'serial_number': self.serial_number,
            'not_before': self.not_before.isoformat(),
            'not_after': self.not_after.isoformat(),
            'public_key': self.public_key,
            'signature': self.signature,
            'cert_type': self.cert_type
        }

    @staticmethod
    def from_dict(data):
        return Certificate(
            subject=data['subject'],
            issuer=data['issuer'],
            serial_number=data['serial_number'],
            not_before=datetime.fromisoformat(data['not_before']),
            not_after=datetime.fromisoformat(data['not_after']),
            public_key=data['public_key'],
            signature=data['signature'],
            cert_type=data.get('cert_type', 'encryption'),
            version=data.get('version', 'v1'),
            sig_alg=data.get('sig_alg', 'sm3')
        )


class SimpleCA:
    def __init__(self, ca_private_key, ca_public_key):
        self.ca_private_key = ca_private_key
        self.ca_public_key = ca_public_key
        self.conn = sqlite3.connect(DATABASE)
        self.create_table()
        self.update_table()
        self.key_file = "encryption.key"
        if not os.path.exists(self.key_file):
            self.key = Fernet.generate_key()
            with open(self.key_file, 'wb') as f:
                f.write(self.key)
        else:
            with open(self.key_file, 'rb') as f:
                self.key = f.read()
        self.cipher = Fernet(self.key)

    def create_table(self):
        with self.conn:
            self.conn.execute('''CREATE TABLE IF NOT EXISTS certificates (
                subject TEXT,
                issuer TEXT,
                serial_number INTEGER PRIMARY KEY,
                not_before TEXT,
                not_after TEXT,
                public_key TEXT,
                signature TEXT,
                version TEXT,
                sig_alg TEXT,
                cert_type TEXT)''')
            self.conn.execute('''CREATE TABLE IF NOT EXISTS revoked (
                subject TEXT,
                issuer TEXT,
                serial_number INTEGER PRIMARY KEY,
                not_before TEXT,
                not_after TEXT,
                public_key TEXT,
                signature TEXT,
                version TEXT,
                sig_alg TEXT,
                cert_type TEXT)''')

    def update_table(self):
        try:
            self.conn.execute("ALTER TABLE certificates ADD COLUMN cert_type TEXT")
        except sqlite3.OperationalError:
            pass

    def create_certificate(self, subject, valid_days, cert_type):
        serial_number = self.get_next_serial_number()
        not_before = datetime.now()
        not_after = not_before + timedelta(days=valid_days)
        private_key = func.random_hex(32)
        sm2_crypt = CryptSM2(private_key=private_key, public_key='')
        public_key = sm2_crypt._kg(int(private_key, 16), sm2_crypt.ecc_table['g'])
        cert_data = f"{subject}{serial_number}{not_before}{not_after}{public_key}"
        K = func.random_hex(16)
        signature = sm2_crypt.sign(cert_data.encode(), K)
        cert = Certificate(subject, 'CA', serial_number, not_before, not_after, public_key, signature, cert_type)
        self.save_certificate(cert)
        # 加密私钥并保存
        encrypted_private_key = self.cipher.encrypt(private_key.encode())
        self.save_private_key(subject, encrypted_private_key)
        return cert, encrypted_private_key

    def save_private_key(self, subject, encrypted_private_key):
        key_folder = "keys"
        os.makedirs(key_folder, exist_ok=True)
        key_path = os.path.join(key_folder, f"{subject}_private_key.enc")
        with open(key_path, 'wb') as f:
            f.write(encrypted_private_key)

    # 解密私钥
    def decrypt_private_key(self, subject):
        key_folder = "keys"
        key_path = os.path.join(key_folder, f"{subject}_private_key.enc")
        if not os.path.exists(key_path):
            print(f"Private key file for {subject} not found.")
            return None

        with open(key_path, 'rb') as f:
            encrypted_private_key = f.read()

        try:
            decrypted_private_key = self.cipher.decrypt(encrypted_private_key).decode()
            return decrypted_private_key
        except Exception as e:
            print(f"Error decrypting private key for {subject}: {e}")
            return None

    def get_next_serial_number(self):
        result = self.conn.execute("SELECT MAX(serial_number) FROM certificates").fetchone()
        return (result[0] or 0) + 1

    def save_certificate(self, cert):
        with self.conn:
            self.conn.execute(
                '''INSERT INTO certificates (subject, issuer, serial_number, not_before, not_after, public_key, signature, version, sig_alg, cert_type)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                (cert.subject, cert.issuer, cert.serial_number, cert.not_before.isoformat(),
                 cert.not_after.isoformat(), cert.public_key, cert.signature, cert.version, cert.sig_alg, cert.cert_type)
            )

    def revoke_certificate(self, serial_number):
        with self.conn:
            cert_row = self.conn.execute("SELECT * FROM certificates WHERE serial_number = ?",
                                         (serial_number,)).fetchone()
            if cert_row:
                self.conn.execute('''INSERT INTO revoked (subject, issuer, serial_number, not_before, not_after, public_key, signature, version, sig_alg, cert_type)
                                     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                                  cert_row)
                self.conn.execute("DELETE FROM certificates WHERE serial_number = ?", (serial_number,))

    def update_certificate(self, serial_number, new_valid_days):
        with self.conn:
            cert_row = self.conn.execute("SELECT * FROM certificates WHERE serial_number = ?", (serial_number,)).fetchone()
            if cert_row:
                not_after = datetime.now() + timedelta(days=new_valid_days)
                cert_data = f"{cert_row[0]}{cert_row[2]}{cert_row[3]}{not_after.isoformat()}{cert_row[5]}"
                sm2_crypt = CryptSM2(private_key=self.ca_private_key, public_key=self.ca_public_key)
                K = func.random_hex(16)
                new_signature = sm2_crypt.sign(cert_data.encode(), K)
                self.conn.execute(
                    '''UPDATE certificates
                    SET not_after = ?, signature = ?
                    WHERE serial_number = ?''',
                    (not_after.isoformat(), new_signature, serial_number)
                )

    def archive_certificate(self, cert):
        archive_folder = "archive"
        os.makedirs(archive_folder, exist_ok=True)
        cert_filename = f"{cert.serial_number}.json"
        cert_path = os.path.join(archive_folder, cert_filename)

        # 检查是否有同名文件
        if os.path.exists(cert_path):
            base, ext = os.path.splitext(cert_filename)
            i = 1
            while os.path.exists(cert_path):
                cert_filename = f"{base}({i}){ext}"
                cert_path = os.path.join(archive_folder, cert_filename)
                i += 1

        with open(cert_path, 'w') as f:
            json.dump(cert.to_dict(), f)
        with open(cert_path, 'rb') as f:
            encrypted_data = self.cipher.encrypt(f.read())
        with open(cert_path, 'wb') as f:
            f.write(encrypted_data)

        # 删除证书或吊销数据库中的记录
        with self.conn:
            self.conn.execute("DELETE FROM certificates WHERE serial_number = ?", (cert.serial_number,))
            self.conn.execute("DELETE FROM revoked WHERE serial_number = ?", (cert.serial_number,))

    def load_certificate(self, serial_number):
        cursor = self.conn.execute("SELECT * FROM certificates WHERE serial_number = ?", (serial_number,))
        row = cursor.fetchone()
        if row:
            return Certificate(
                subject=row[0],
                issuer=row[1],
                serial_number=row[2],
                not_before=datetime.fromisoformat(row[3]),
                not_after=datetime.fromisoformat(row[4]),
                public_key=row[5],
                signature=row[6],
                version=row[7],
                sig_alg=row[8],
                cert_type=row[9]
            )
        return None

    def load_revoked_certificate(self, serial_number):
        cursor = self.conn.execute("SELECT * FROM revoked WHERE serial_number = ?", (serial_number,))
        row = cursor.fetchone()
        if row:
            return Certificate(
                subject=row[0],
                issuer=row[1],
                serial_number=row[2],
                not_before=datetime.fromisoformat(row[3]),
                not_after=datetime.fromisoformat(row[4]),
                public_key=row[5],
                signature=row[6],
                version=row[7],
                sig_alg=row[8],
                cert_type=row[9]
            )
        return None

# 操作函数
def initialize_ca():
    ca_private_key = func.random_hex(32)
    sm2_crypt = CryptSM2(private_key=ca_private_key, public_key='')
    ca_public_key = sm2_crypt._kg(int(ca_private_key, 16), sm2_crypt.ecc_table['g'])
    return SimpleCA(ca_private_key, ca_public_key)


def create_certificate(ca, subject, valid_days, cert_type):
    cert, private_key = ca.create_certificate(subject, valid_days, cert_type)
    print(f"{cert_type.capitalize()} Certificate created successfully:\n", cert.to_dict())
    return cert, private_key


def revoke_certificate(ca, serial_number):
    ca.revoke_certificate(serial_number)
    print(f"Certificate with serial number {serial_number} has been revoked.")


def update_certificate(ca, serial_number, new_valid_days):
    ca.update_certificate(serial_number, new_valid_days)
    print(f"Certificate with serial number {serial_number} has been updated.")


def archive_certificate(ca, cert):
    ca.archive_certificate(cert)
    print(f"Certificate with serial number {cert.serial_number} has been archived.")


def load_certificate(ca, serial_number):
    cert = ca.load_certificate(serial_number)
    if cert:
        print(f"Loaded Certificate:\n", cert.to_dict())
    else:
        print(f"Certificate with serial number {serial_number} not found.")
    return cert

def load_revoked_certificate(ca, serial_number):
    cert = ca.load_revoked_certificate(serial_number)
    if cert:
        print(f"Loaded Revoked Certificate:\n", cert.to_dict())
    else:
        print(f"Revoked certificate with serial number {serial_number} not found.")
    return cert


def main():
    ca = initialize_ca()
    while True:
        print("\n请选择你要进行的操作：")
        print("1. 创建证书")
        print("2. 吊销证书")
        print("3. 更新证书")
        print("4. 归档证书")
        print("5. 加载证书")
        print("6. 加载吊销的证书")
        print("7. 退出")

        choice = input("请输入操作序号: ")
        if choice == '1':
            subject = input("请输入证书主题: ")
            valid_days = int(input("请输入证书有效期（天）: "))
            cert_type = input("请输入证书类型（encryption/signature）: ")
            create_certificate(ca, subject, valid_days, cert_type)
        elif choice == '2':
            serial_number = int(input("请输入要吊销的证书序列号: "))
            revoke_certificate(ca, serial_number)
        elif choice == '3':
            serial_number = int(input("请输入要更新的证书序列号: "))
            new_valid_days = int(input("请输入新的有效期（天）: "))
            update_certificate(ca, serial_number, new_valid_days)
        elif choice == '4':
            serial_number = int(input("请输入要归档的证书序列号: "))
            cert = load_certificate(ca, serial_number)
            if not cert:
                cert = load_revoked_certificate(ca, serial_number)
            if cert:
                archive_certificate(ca, cert)
            else:
                print(f"证书序列号 {serial_number} 未找到。")
        elif choice == '5':
            serial_number = int(input("请输入要加载的证书序列号: "))
            load_certificate(ca, serial_number)
        elif choice == '6':
            serial_number = int(input("请输入要加载的吊销证书序列号: "))
            load_revoked_certificate(ca, serial_number)
        elif choice == '7':
            break
        else:
            print("无效的操作序号，请重试。")

def archive_certificate(ca, cert):
    archive_folder = "archive"
    os.makedirs(archive_folder, exist_ok=True)
    cert_filename = f"{cert.serial_number}.json"
    cert_path = os.path.join(archive_folder, f"{cert.serial_number}.json")

    # 检查是否有同名文件
    if os.path.exists(cert_path):
        base, ext = os.path.splitext(cert_filename)
        i = 1
        while os.path.exists(cert_path):
            cert_filename = f"{base}({i}){ext}"
            cert_path = os.path.join(archive_folder, cert_filename)
            i += 1

    with open(cert_path, 'w') as f:
        json.dump(cert.to_dict(), f)
    with open(cert_path, 'rb') as f:
        encrypted_data = ca.cipher.encrypt(f.read())
    with open(cert_path, 'wb') as f:
        f.write(encrypted_data)

    # 删除证书或吊销数据库中的记录
    with ca.conn:
        ca.conn.execute("DELETE FROM certificates WHERE serial_number = ?", (cert.serial_number,))
        ca.conn.execute("DELETE FROM revoked WHERE serial_number = ?", (cert.serial_number,))

if __name__ == "__main__":
    main()

